from robot.api import logger
from robotlibcore import keyword

from Crypto.Cipher import AES
import base64

class AesLib(object):
    
    def add_to_32(self, value):
        while len(value) % 32 != 0:
            value += b'\x00'
        
        return value
    
    def cut_value(self, org_str):
        org_bytes = str.encode(org_str)
        n = int(len(org_bytes) / 32)
        i = 0
        new_bytes = b''
        while n >= 1:
            i = i + 1
            new_byte = org_bytes[(i-1)*32:32*i-1]
            new_bytes += new_byte
            n = n - 1
        if len(org_bytes) % 32 == 0:                   
            all_bytes = org_bytes
        elif len(org_bytes) % 32 != 0 and n>1:         
            all_bytes = new_bytes + self.add_to_32 (org_bytes[i*32:])
        else:
            all_bytes = self.add_to_32 (org_bytes) 
        return all_bytes

    @keyword('AES Decrypt')
    def AES_decrypt(self, secret_str, key="aaaaaaaaaabbbbbb"):

        aes = AES.new(self.cut_value(key), AES.MODE_CBC)
        base64_decrypted = base64.decodebytes(secret_str.encode(encoding='utf-8'))
        decrypted_text = str(aes.decrypt(base64_decrypted), encoding='utf-8').replace('\0', '')

        return decrypted_text